﻿///  from http://perso.univ-lemans.fr/~berger/Afsd56/sample_Levenb.cpp
#include <iostream>
#include <vector>

#include <opencv2/imgcodecs.hpp>
#include <opencv2/videoio/videoio.hpp>
#include <opencv2/highgui/highgui.hpp>

using namespace std;
using namespace cv;

class LMSolver : public Algorithm
{
public:
    class Callback
    {
    public:
        virtual ~Callback() {}
        virtual bool compute(InputArray param, OutputArray err, OutputArray J) const = 0;
    };

    virtual void setCallback(const Ptr<LMSolver::Callback>& cb) = 0;
    virtual int run(InputOutputArray _param0) const = 0;
};

class LMSolverImpl : public LMSolver
{
public:
    LMSolverImpl() : maxIters(100) { init(); }
    LMSolverImpl(const Ptr<LMSolver::Callback>& _cb, int _maxIters) : cb(_cb), maxIters(_maxIters) { init(); }

    void init()
    {
        epsx = epsf = FLT_EPSILON;
        printInterval = 1;
    }

    void setCallback(const Ptr<LMSolver::Callback>& _cb) { cb = _cb; }

    int run(InputOutputArray _param0) const
    {
        Mat param0 = _param0.getMat(), x, xd, r, rd, J, A, Ap, v, temp_d, d;
        int ptype = param0.type();

        CV_Assert( (param0.cols == 1 || param0.rows == 1) && (ptype == CV_32F || ptype == CV_64F));
        CV_Assert( cb );

        int lx = param0.rows + param0.cols - 1;
        param0.convertTo(x, CV_64F);

        if( x.cols != 1 )
            transpose(x, x);
        // x 变成列向量

        // x列向量参数，r 残差/误差  J 一阶偏导数矩阵
        if( !cb->compute(x, r, J) )
            return -1;
        // 2-范数的平方 || r||^2_2
        double S = norm(r, NORM_L2SQR);  // 目标S -> 0
        int nfJ = 2;

        // A = J(t)J
        mulTransposed(J, A, true);
        // v = J(t)*r = J(t)*f
        gemm(J, r, 1, noArray(), 0, v, GEMM_1_T);

        Mat D = A.diag().clone();

        const double Rlo = 0.25, Rhi = 0.75; // 信赖域 两个判断条件 0.25 < p < 0.75
        double lambda = 1, lc = 0.75;
        int i, iter = 0;

        if( printInterval != 0 )
        {
            printf("************************************************************************************\n");
            printf("\titr\tnfJ\t\tSUM(r^2)\t\tx\t\tdx\t\tl\t\tlc\n");
            printf("************************************************************************************\n");
        }
        for( ;; )
        {
            CV_Assert( A.type() == CV_64F && A.rows == lx );
            A.copyTo(Ap);
            // J^T J + uI => Ap
            for ( i = 0; i < lx; i++) {
                Ap.at<double>(i, i) += lambda * D.at<double>(i);
            }

            // (J^T J + uI)*h = -J^T*f(x)   => Ap * d = v
            // 这里d = -h
            solve(Ap, v, d, DECOMP_EIG);
            subtract(x, d, xd);
            // 此时 xd = x+h = x-d 为 k+1 一次迭代结果
            if( !cb->compute(xd, rd, noArray()) )
                return - 1;
            nfJ++;
            // Sd = F(X+h)  S = F(X)
            double Sd = norm(rd, NORM_L2SQR);
            // temp_d = -A^T d + 2 v^T = J.t*h+2J(t)*f
            gemm(A, d, -1, v, 2, temp_d);
            double dS = d.dot(temp_d);
            // R = (F(x) - F(x+h))/(L(0) - L(h))
            double R = (S - Sd)/(fabs(dS) > DBL_EPSILON ? dS : 1);

            if( R > Rhi )
            {
                lambda *= 0.5;
                if( lambda < lc )
                    lambda = 0;
            }
            else if( R < Rlo )
            {
                // find new nu if R too low
                double t = d.dot(v);
                double nu = (Sd - S)/(fabs(t) > DBL_EPSILON ? t : 1) + 2;
                nu = std::min(std::max(nu, 2.), 10.);
                if( lambda == 0 )
                {
                    invert(A, Ap, DECOMP_EIG);
                    double maxval = DBL_EPSILON;
                    for( i = 0; i < lx; i++ )
                        maxval = std::max(maxval, std::abs(Ap.at<double>(i,i)));
                    lambda = lc = 1./maxval;
                    nu *= 0.5;
                }
                lambda *= nu;
            }

            if( Sd < S )
            {
                nfJ++;
                S = Sd;
                std::swap(x, xd);
                if( !cb->compute(x, r, J) )
                    return -1;
                mulTransposed(J, A, true);
                gemm(J, r, 1, noArray(), 0, v, GEMM_1_T);
            }

            iter++;
            bool proceed = iter < maxIters && norm(d, NORM_INF) >= epsx && norm(r, NORM_INF) >= epsf;

            if( printInterval != 0 && (iter % printInterval == 0 || iter == 1 || !proceed) )
            {
                printf("%c%10d %10d %15.4e %16.4e %17.4e %16.4e %17.4e\n",
                       (proceed ? ' ' : '*'), iter, nfJ, S, x.at<double>(0), d.at<double>(0), lambda, lc);
            }

            if(!proceed)
                break;
        }

        if( param0.size != x.size )
            transpose(x, x);

        x.convertTo(param0, ptype);
        if( iter == maxIters )
            iter = -iter;

        return iter;
    }


    Ptr<LMSolver::Callback> cb;
    double epsx;
    double epsf;
    int maxIters;
    int printInterval;
};


Ptr<LMSolver> createLMSolver(const Ptr<LMSolver::Callback>& cb, int maxIters)
{
    return makePtr<LMSolverImpl>(cb, maxIters);
}

class ModeleCallback : public LMSolver::Callback
{
public:
    ModeleCallback(InputArray _src, InputArray _dst)
    {
        src = _src.getMat();
        dst = _dst.getMat();
        cout << "src:" << src << endl;
        cout << "dst:" << dst << endl;
    }

    bool compute(InputArray _param, OutputArray _err, OutputArray _Jac) const
    {
        int i, count = src.cols;     // 传入的src是行向量
        Mat param = _param.getMat(); // 是列向量
        _err.create(count, 1, CV_64F);
        Mat err = _err.getMat(), J;
        if( _Jac.needed())
        {
            _Jac.create(count, param.rows, CV_64F); // 行 = src.cols 列 = param.rows
            J = _Jac.getMat();
            CV_Assert( J.isContinuous()  );
        }

        // m = x * h
        const double* x = src.ptr<double>();
        const double* m = dst.ptr<double>();
        const double* h = param.ptr<double>();

        double* errptr = err.ptr<double>();
        double* Jptr = J.data ? J.ptr<double>() : 0;

        for( i = 0; i < count; i++ )
        {
            double xi = 1/ (h[0] * x[i] * x[i] + 1) + sin(h[1] * x[i] + h[2]);
            errptr[i] = xi - m[i];

            // J 偏导矩阵
            if(Jptr) {
                // 一行，三列
                // 对 a 偏导， 也就是 h[0]
                Jptr[0] =  -x[i] * x[i]/((h[0] * x[i] * x[i] + 1)*(h[0] * x[i] * x[i] + 1));
                // 对 b 偏导
                Jptr[1] = x[i] * cos(h[1] * x[i] + h[2]);
                // 对 c 偏导
                Jptr[2] = cos(h[1] * x[i] + h[2]);
                Jptr += 3;
            }
        }

        return true;
    }

    Mat src, dst;
};

int main(int argc, char *argv[])
{
    vector<double> param = {1,6,2}; // 1/(ax^2+1)+sin(bx+c)
    vector<double> paramIni = {1.5,4,1}; // 1/(ax^2+1)+sin(bx+c)

    vector<double> x = {0,.1,.2,.3,.4,.5,.6,.7,.8,.9,1.0}; // x data;
    vector<double> y(x.size());

    for (int i = 0; i<x.size();i++)
        y[i] =  1/(param[0] * x[i] * x[i] + 1) + sin(param[1] * x[i] + param[2]);

    for (int i = 0; i<x.size();i++)
        cout << x[i] << "\t"<<y[i]<<endl;

    createLMSolver(makePtr<ModeleCallback>(x, y), 100)->run(paramIni);
    cout << "Param result\n";
    for (int i = 0; i<paramIni.size();i++)
        cout << paramIni[i] <<endl;
    cout << "Error\n";
    for (int i = 0; i<x.size();i++)
        cout<<y[i]-( 1/(paramIni[0] * x[i] * x[i] + 1) + sin(paramIni[1] * x[i] + paramIni[2]))<<endl;
    return 0;
}
